from pydantic import BaseModel, Field
from langchain_deepseek import ChatDeepSeek
from langchain.prompts import PromptTemplate
# 使用JsonOutputParser解析输出
from langchain_core.output_parsers import JsonOutputParser 

class Date(BaseModel):
    year: int = Field(description="Year")
    month: int = Field(description="Month")
    day: int = Field(description="Day")
    era: str = Field(description="BC or AD")

parser = JsonOutputParser(pydantic_object=Date)

prompt = PromptTemplate(
    template="提取用户输入中的日期。\n用户输入:{query}\n{format_instructions}",
    input_variables=["query"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

input_prompt = prompt.format_prompt(query="2024年十二月23日天气晴...")

llm = ChatDeepSeek(model="deepseek-chat", temperature=0)
output = llm.invoke(input_prompt)
print("原始输出:\n"+output.content)

print("\n解析后:")
result = parser.invoke(output)
print(result)
print(type(result))